#include "board.h"
#include "z80.h"
#include "z80_sd.h"

//----------------------------------------------------------------------------
// SD Registers
//----------------------------------------------------------------------------
static uint8_t  Reg_Command = 0;        // Command register
static uint32_t Reg_BlockAddress = 0;   // Block address register 32-bit
static uint8_t  Reg_Status = 0;         // Status register
static uint16_t DataBufferIndex = 0;    // Data access index 
static uint8_t  DataBuffer[512] = {0};  // Data buffer for R/W

//----------------------------------------------------------------------------
// SD Card SPI Command Definitions
//----------------------------------------------------------------------------
// Basic commands (CMD0-CMD59)
#define CMD0    0x40  // GO_IDLE_STATE: reset card, enter SPI mode
#define CMD1    0x41  // SEND_OP_COND (MMC only)
#define CMD8    0x48  // SEND_IF_COND: check voltage range, V2 cards
#define CMD9    0x49  // SEND_CSD: read CSD register
#define CMD10   0x4A  // SEND_CID: read CID register
#define CMD12   0x4C  // STOP_TRANSMISSION
#define CMD13   0x4D  // SEND_STATUS: read card status
#define CMD16   0x50  // SET_BLOCKLEN: set read/write block length
#define CMD17   0x51  // READ_SINGLE_BLOCK
#define CMD18   0x52  // READ_MULTIPLE_BLOCK
#define CMD23   0x57  // SET_BLOCK_COUNT (MMC only)
#define CMD24   0x58  // WRITE_BLOCK (write single block)
#define CMD25   0x59  // WRITE_MULTIPLE_BLOCK
#define CMD55   0x77  // APP_CMD: next command is application-specific
#define CMD58   0x7A  // READ_OCR: read OCR register
#define CMD59   0x7B  // CRC_ON_OFF: enable/disable CRC checking

// Application specific commands (ACMD)
#define ACMD41  0x69  // SD_SEND_OP_COND: initialization (must be after CMD55)
#define ACMD13  0x4D  // SD_STATUS: read SD status register
#define ACMD23  0x57  // SET_WR_BLK_ERASE_COUNT (SD only)

// Data tokens
#define TOKEN_START_BLOCK     0xFE  // Single block read/write start token
#define TOKEN_MULTI_WRITE     0xFC  // Start token for multiple block write
#define TOKEN_STOP_TRAN       0xFD  // Stop transmission token for multiple write

//----------------------------------------------------------------------------
// SPI Delay Function
//----------------------------------------------------------------------------
// Simple short delay to control SPI timing
void SPI_Delay(void)
{
    for (volatile int i = 0; i < 0; i++);
}

//----------------------------------------------------------------------------
// SPI Transfer Function
//----------------------------------------------------------------------------
// Send and receive one byte via software SPI (bit-banging)
// SPI Mode 0 (CPOL=0, CPHA=0): clock idle low, sample on rising edge
uint8_t SPI_Transfer(uint8_t data)
{
    *(volatile int*)REG_SPI0_DATA = data;
    return *(volatile int*)REG_SPI0_DATA & 0xFF;
}

//----------------------------------------------------------------------------
// SPI Read & Write
//----------------------------------------------------------------------------
// Read one byte from SPI (sending 0xFF as dummy data)
uint8_t SPI_ReadByte(void)
{
    return SPI_Transfer(0xFF);
}

// Write one byte to SPI
void SPI_WriteByte(uint8_t data)
{
    SPI_Transfer(data);
}

//----------------------------------------------------------------------------
// SD Card Command Interface
//----------------------------------------------------------------------------
// Send one command to SD card and return R1 response (1 byte)
uint8_t SD_SendCommand(uint8_t cmd, uint32_t arg, uint8_t crc)
{
    uint8_t resp, i;

    // Ensure CS high before asserting low
    SPI_ReadByte();

    // Send command packet
    SPI_WriteByte(cmd);
    SPI_WriteByte(arg >> 24);
    SPI_WriteByte(arg >> 16);
    SPI_WriteByte(arg >> 8);
    SPI_WriteByte(arg);
    SPI_WriteByte(crc);

    // Wait for response (response bit7=0)
    for (i = 0; i < 8; i++) {
        resp = SPI_ReadByte();
        if (!(resp & 0x80)) break;
    }
    return resp;
}

//----------------------------------------------------------------------------
// SD Card Initialization
//----------------------------------------------------------------------------
// Initialize SD card in SPI mode
uint8_t SD_Initialize(void)
{
    // Set SPI speed 400KHz
    *(volatile int*)REG_SPI0_CLK_DIV = SPI_CLK_DIV_0;

    uint8_t i, resp;

    // Send 80 clock cycles with CS high before initialization
    for (i = 0; i < 10; i++)
        SPI_ReadByte();

    // CMD0: Enter Idle state
    resp = SD_SendCommand(CMD0, 0, 0x95);
    if (resp != 0x01) return 1;

    // CMD8: Check card version and voltage range
    resp = SD_SendCommand(CMD8, 0x000001AA, 0x87);
    if (resp == 0x01) {
        // Version 2.x card (supports CMD8)
        uint32_t timeout = 0x0FFFFFFF;
        do {
            SD_SendCommand(CMD55, 0, 0x01);
            resp = SD_SendCommand(ACMD41, 0x40000000, 0x01);
        } while (resp != 0x00 && --timeout);

        if(timeout == 0) {
            return 2;
        }
    } else {
        // Version 1.x card (does not support CMD8)
        do {
            SD_SendCommand(CMD55, 0, 0x01);
            resp = SD_SendCommand(ACMD41, 0, 0x01);
        } while (resp != 0x00);
    }

    // CMD58: Read OCR (operating conditions)
    SD_SendCommand(CMD58, 0, 0x01);
    for (i = 0; i < 4; i++)
        SPI_ReadByte();

    SPI_ReadByte();

    // Set SPI speed 20MHz
    *(volatile int*)REG_SPI0_CLK_DIV = SPI_CLK_DIV_1;

    return 0;
}

//----------------------------------------------------------------------------
// SD Read Single Block
//----------------------------------------------------------------------------
// Read one 512-byte data block from SD card
uint8_t SD_ReadBlock(uint32_t blockAddr, uint8_t* buf)
{
    uint8_t token;

    // CMD17: Read single block command
    if (SD_SendCommand(CMD17, blockAddr, 0x01))
        return 1;

    // Wait for data token (0xFE)
    for (uint32_t i = 0; i < 0x0FFFFFFF; i++) {
        token = SPI_ReadByte();
        if (token == TOKEN_START_BLOCK)
            break;
    }

    if (token != TOKEN_START_BLOCK)
        return 2;

    // Read 512 bytes of data
    for (uint16_t i = 0; i < 512; i++)
        buf[i] = SPI_ReadByte();

    // Read and discard CRC
    SPI_ReadByte();
    SPI_ReadByte();

    SPI_ReadByte();
    return 0;
}

//----------------------------------------------------------------------------
// SD Write Single Block
//----------------------------------------------------------------------------
// Write a single 512-byte block to SD card
uint8_t SD_WriteBlock(uint32_t blockAddr, const uint8_t* buf)
{
    //printf("SD Write Block:%d\r\n\r\n", blockAddr);
    uint8_t resp;
    uint16_t i;

    // Send WRITE_SINGLE_BLOCK command (CMD24)
    resp = SD_SendCommand(CMD24, blockAddr, 0x01);
    if (resp != 0x00) {
        SPI_ReadByte();
        return 1; // Command failed
    }

    // Send one dummy byte before data token
    SPI_WriteByte(0xFF);

    // Send start token (0xFE)
    SPI_WriteByte(TOKEN_START_BLOCK);

    // Write 512-byte data block
    for (i = 0; i < 512; i++) {
        SPI_WriteByte(buf[i]);
    }

    // Send dummy CRC
    SPI_WriteByte(0xFF);
    SPI_WriteByte(0xFF);

    // Read data response token
    resp = SPI_ReadByte();
    if ((resp & 0x1F) != 0x05) { // 0x05 = data accepted
        SPI_ReadByte();
        return 2; // Data rejected
    }

    // Wait until card is not busy
    while (SPI_ReadByte() == 0x00);

    SPI_ReadByte();
    return 0; // Success
}

// Initialize GPIO, then initialize and test SD card
void SD_Init()
{
    // Enable GPIO4 peripheral clock
    SYS_EnableAPBClock(APB_MASK_GPIO4);

    // Set GPIO4_1,2,3 as output (CS, SCK, MOSI)
    GPIO_SetOutput(GPIO4, GPIO_BIT1|GPIO_BIT2|GPIO_BIT3);

    // Set GPIO4_4 as input (MISO)
    GPIO_SetInput(GPIO4, GPIO_BIT4);

    // Set initial output state to high
    GPIO_SetHigh(GPIO4, GPIO_BIT1|GPIO_BIT2|GPIO_BIT3);

    // Initialize SD card
    SD_Initialize();
}

void SD_Read(uint8_t addr, uint8_t* out)
{
    // Handle data register read
    if(addr == Z80_SD_DATA) {
        if(DataBufferIndex >= 512) {
            DataBufferIndex = 0;
        }
        *out = DataBuffer[DataBufferIndex++];
    }
    // Handle status register read
    else if(addr == Z80_SD_STATUS) {
        *out = Reg_Status;
    }
}

void SD_Write(uint8_t addr, const uint8_t* in)
{
    // Handle data register write
    if(addr == Z80_SD_DATA) {
        DataBuffer[DataBufferIndex++] = *in;

        // Write to the SD card when the buffer reaches 512 bytes
        if(DataBufferIndex >= 512) {
            DataBufferIndex = 0;
            uint8_t val = SD_WriteBlock(Reg_BlockAddress, DataBuffer);
            Reg_Status = (val == 0 ? ERR_OK : ERR_WRITE_BLOCK);

            //printf("Write Block Status: %d, Data: %s\r\n\r\n", Reg_Status, DataBuffer);
        }
    }
    // Handle block register write
    else if(addr == Z80_SD_BLOCK_ADDR0) {
        Reg_BlockAddress &= 0xFFFFFF00;
        Reg_BlockAddress |= *in;
    }
    else if(addr == Z80_SD_BLOCK_ADDR1) {
        Reg_BlockAddress &= 0xFFFF00FF;
        Reg_BlockAddress |= (((uint32_t)*in) << 8) & 0x0000FF00;
    }
    else if(addr == Z80_SD_BLOCK_ADDR2) {
        Reg_BlockAddress &= 0xFF00FFFF;
        Reg_BlockAddress |= (((uint32_t)*in) << 16) & 0x00FF0000;
    } else if(addr == Z80_SD_BLOCK_ADDR3) {
        Reg_BlockAddress &= 0x00FFFFFF;
        Reg_BlockAddress |= (((uint32_t)*in) << 24) & 0xFF000000;
    }
    // Handle command register write
    else if(addr == Z80_SD_COMMAND) {
        if(*in == COMMAND_READ_BLOCK) {
            // Read SD block
            DataBufferIndex = 0;
            uint8_t val = SD_ReadBlock(Reg_BlockAddress, DataBuffer);
            Reg_Status = (val == 0 ? ERR_OK : ERR_READ_BLOCK);


            //printf("Read Block Addr: %d, Status: %d\r\n", Reg_BlockAddress, Reg_Status);
        } else if(*in == COMMAND_WRITE_BLOCK) {
            // Clear data buffer
            DataBufferIndex = 0;

            //printf("Write Block Address, Block: %d\r\n", Reg_BlockAddress);
        }
    }
}